#pragma once

#include "ISpacePartition.h"

class KdTreeNode;
class KdTree;
class NaiveKdTree;


struct CompactKdTreeNode;



class KdTreeNode
{
	friend class KdTree;
	friend class NaiveKdTree;
public: 
	KdTreeNode();
	virtual ~KdTreeNode();

	inline bool IsLeaf() const;
	inline int GetAxis() const;
	inline float GetSplit() const;

	inline KdTreeNode* GetFrontChild() const;
	inline KdTreeNode* GetBackChild() const;

	void TraceRay(const KdTree* pTree, RAY* pRay, TraceType traceType, ITraceResult* pResult) const;
	void GetLocalBoundingBox(const KdTree* pTree, AABB& bb);
	size_t GetTriangleCount() const;

	size_t GetNumberOfNode() const;

	void BuildCompact(
		KdTree* pTree, 
		CompactKdTreeNode* pCurrentNode,
		unsigned int& iAvailableNode,
		unsigned int& iAvailableLeafData);

protected:
	void CleanMembers();

protected:
	int mAxis;		// 0 - X, 1 - Y, 2 - Z
	float mSplit;
	bool mIsLeaf;

	// For non-leaf node
	KdTreeNode* mpFrontChild;
	KdTreeNode* mpBackChild;

	// For leaf node
	vector<int> mTriangleList;

	// Bounding box
	AABB mAABB;
};

bool KdTreeNode::IsLeaf() const
{
	return mIsLeaf;
}

int KdTreeNode::GetAxis() const
{
	return mAxis;
}

float KdTreeNode::GetSplit() const
{
	return mSplit;
}

KdTreeNode* KdTreeNode::GetFrontChild() const
{
	return mpFrontChild;
}

KdTreeNode* KdTreeNode::GetBackChild() const
{
	return mpBackChild;
}




class KdTree: public ISpacePartition
{
	friend class KdTreeNode;
	friend struct CompactKdTreeNode;
	
public:
	KdTree(const vector<TCRTMesh*>& Meshes);
	virtual ~KdTree();

	virtual bool Build();
	virtual bool IsReady() const;

	virtual void TraceRay(RAY* pRay, TraceType traceType, ITraceResult* pResult);

protected:
	virtual void CleanMembers();

	void BuildSubTree(KdTreeNode* pRoot, AABB& bb, size_t iLevel);

	virtual bool StopSplitting(KdTreeNode* pRoot, size_t iLevel) = 0;
	virtual bool FindSplitPlane(KdTreeNode* pRoot, size_t iLevel, int& iAxis, float& fPosition) = 0;

	// Stop splittig criteria
	virtual size_t GetMaximumDepth() const = 0;
	virtual size_t GetMaximumTrianglePerLeaf() const = 0;

	virtual void ReportStats() const;

protected:
	vector<TCRTMesh*> mRefMeshes;
	bool mTreeReady;

	// Vertexs from different meshes are merged into a single list. The number of vertices can not exceed 2G.
	Vertex* mpVertexList;	
	Triangle* mpTriangleList;
	// The same length of mpTriangleList. 
	// mpTriangleList[i].vertexIndex[j] - mpVertexOffsetList[i] yields the actual vertex index into a TCRTMesh
	int* mpVertexOffsetList;	
	// The same length of mpTriangleList. 
	// Contain a binary flag which indicated whether the indexed triangle has been performed 
	// triangle/ray intersection testing for current ray. 
	bool* mpTriangleVisitedFlag;

	size_t mVertexCount;
	size_t mTriangleCount;

	KdTreeNode* mpRoot;

	size_t mNodeNum;
	size_t mLeafTriangleNum;
	CompactKdTreeNode* mpCompactRoot;
	unsigned int* mpLeafTriangleData;

	size_t mMaximumDepth;
};

// of size 32 bytes
struct CompactKdTreeNode
{
	// bit31 (most important bit) - is leaf?
	// bit 30 & 29: axis (0-x, 1-y, 2-z) for inner, and reserved for leaf
	// bit 28 - 0: 
	//		point into the tree's mpLeafTriangleData if leaf. 
	//		point into the tree's mpNodeList if inner
	unsigned int mMultiUsageFlag;

	union
	{
		float mSplit;				// For Inner only
		unsigned int mTriangleNum;	// For Leaf only
	} mU;

	AABB mAABB;

	void TraceRay(const KdTree* pTree, RAY* pRay, TraceType traceType, ITraceResult* pResult) const;

	inline bool IsLeaf() const;
	inline int Axis() const;
	inline float Split() const;
	inline unsigned int TriangleNum() const;
	inline unsigned int* TriangleList(const KdTree* pTree) const;
	inline CompactKdTreeNode* FrontChild(const KdTree* pTree) const;
	inline CompactKdTreeNode* BackChild(const KdTree* pTree) const;
};

inline bool CompactKdTreeNode::IsLeaf() const
{
	return (mMultiUsageFlag & 0x80000000) != 0;
}

inline int CompactKdTreeNode::Axis() const
{
	return (mMultiUsageFlag & 0x70000000) >> 29;
}

inline float CompactKdTreeNode::Split() const
{
	return mU.mSplit;
}

inline unsigned int CompactKdTreeNode::TriangleNum() const
{
	return mU.mTriangleNum;
}

inline unsigned int* CompactKdTreeNode::TriangleList(const KdTree* pTree) const
{
	unsigned int index = mMultiUsageFlag & 0x1fffffff;
	return pTree->mpLeafTriangleData + index;
}

inline CompactKdTreeNode* CompactKdTreeNode::FrontChild(const KdTree* pTree) const
{
	unsigned int index = mMultiUsageFlag & 0x1fffffff;
	return pTree->mpCompactRoot + index;
}

inline CompactKdTreeNode* CompactKdTreeNode::BackChild(const KdTree* pTree) const
{
	unsigned int index = mMultiUsageFlag & 0x1fffffff;
	return pTree->mpCompactRoot + index + 1;
}


